#!/usr/bin/python
# Intel SR-IOV based network card Virtual Function utility for kvm
# Python 2.7.X
# Shankhadeep Shome
# Version 0.3

# Third party libraries not part of standard Ubuntu install required
# pexpect - http://www.noah.org/wiki/Pexpect
# prettytable - apt-get install python-prettytable
#
# Third party utilities required: virsh, libvirtd, qemu-kvm

import os
import sys
import subprocess
import optparse
import time
import datetime
import prettytable
import re

# Python class VFinfo

class VFinfo:
	def __init__(self):
		self.VFnumber = ''
		self.VFparentdev = ''
		self.VFparentdriver = ''
		self.VFpciid = ''
		self.VFmac = ''
		self.VFvlan = ''
		self.VFtxrate = ''
		self.VFlinkstatus = ''
		self.VFassignedvm = ''
		self.VFflexpp = ''
	
# Input parser

def parse_options():
	usage_string = "usage: %prog [-adcU] [-v vlan] [-r rate] [-m mac] [-M mac] [-f vf_number ] [-p pf_devname] -I [VM instance]"
	version_string = "Version 0.3"
	desc_string = "Intel SR-IOV based network card Virtual Function utility for qemu-kvm"
	parser = optparse.OptionParser(usage=usage_string, version=version_string, description=desc_string)
	parser.add_option("-a", "--Attach", dest="Attach", default=False, action="store_true", help="Attach VF to a running VM")
	parser.add_option("-d", "--Detach", dest="Detach", default=False, action="store_true", help="Detach VF from a running VM")
	parser.add_option("-c", "--Change", dest="Change", default=False, action="store_true", help="Change VF properties")
	parser.add_option("-U", "--Udev", dest="Udev", default=False, action="store_true", help="Dump suggested udev rules for VFs for a PF")
	parser.add_option("-m", "--Mac", dest="Mac", type="string", default=None, help="Change MAC address of a VF")
	parser.add_option("-M", "--Macall", dest="Macall", type="string", default=None, help="Change the MAC address of all the VF of a PF incrementally, keeping the first 5 octets unchanged, running this ignores other arguments")
	parser.add_option("-v", "--Vlan", dest="Vlan", type="string", default=None, help="VLAN # between 0 and 4095, 0 means disable vlan")
	parser.add_option("-r", "--Rate", dest="Rate", type="string", default=None, help="TX rate limit between 0 and 10000 Mbits, 0 means uncap txrate")
	parser.add_option("-f", "--Function", dest="Function", type="string", default=None, help="Virtual Function number")
	parser.add_option("-p", "--Physicaldev", dest="Physicaldev", type="string", default=None, help="Physical device name")
	parser.add_option("-I", "--Instance", dest="Instance", type="string", default=None, help="KVM instance name")
	
	# command line options
	(opts, args) = parser.parse_args()

	if opts.Attach and (opts.Detach or opts.Change or opts.Udev):
		print "Invalid option combination\n"
		parser.print_help()
		exit(-1)
		
	if opts.Detach and (opts.Attach or opts.Change or opts.Udev):
		print "Invalid option combination\n"
		parser.print_help()
		exit(-1)
	
	if opts.Change and (opts.Detach or opts.Attach or opts.Udev):
		print "Invalid option combination\n"
		parser.print_help()
		exit(-1)
	
	if opts.Udev and (opts.Detach or opts.Change or opts.Attach):
		print "Invalid option combination\n"
		parser.print_help()
		exit(-1)
	
	if opts.Attach and not (opts.Function and opts.Physicaldev and opts.Instance):
		print "Attach requires vf number, a physical device and a running kvm instance\n"
		parser.print_help()
		exit(-1)
	
	if opts.Detach and not (opts.Function and opts.Physicaldev and opts.Instance):
		print "Detach requires vf number, a physical device and a running kvm instance\n"
		parser.print_help()
		exit(-1)
	
	if opts.Change and not opts.Physicaldev:
		print "Change requires a physical device\n"
		parser.print_help()
		exit(-1)
	
	if opts.Udev and not opts.Physicaldev:
		print "udev rule dump requires a physical device\n"
		parser.print_help()
		exit(-1)
	
	if (opts.Change and opts.Mac) and not (opts.Function and opts.Physicaldev):
		print "Changing a MAC of a VF requires an VF number and a physical device\n"
		parser.print_help()
		exit(-1)
		
	if (opts.Change and opts.Macall) and not opts.Physicaldev:
		print "Changing MAC of all VFs requires a physical device\n"
		parser.print_help()
		exit(-1)
	
	if (opts.Change and opts.Vlan):
		try: 
			tmp = int(opts.Vlan)
			if (tmp < 0) or (tmp > 4095):
				print "Vlan # out of range\n"
				parser.print_help()
				exit(-1)
		except (ValueError, TypeError):
			print "Vlan # must be an integer\n"
			parser.print_help()
			exit(-1)	
	
	if (opts.Change and opts.Rate):
		try: 
			tmp = int(opts.Rate)
			if (tmp < 0) or (tmp > 10000):
				print "TX Rate out of range\n"
				parser.print_help()
				exit(-1)
		except (ValueError, TypeError):
			print "TX rate must be an integer\n"
			parser.print_help()
			exit(-1)
			
	if (opts.Change and opts.Function):		
		try: 
			tmp = int(opts.Function)
			if (tmp < 0) or (tmp > 63):
				print "VF number out range\n"
				parser.print_help()
				exit(-1)
		except (ValueError, TypeError):
			print "VF number must be an integer\n"
			parser.print_help()
			exit(-1)
	
	if (opts.Change and opts.Mac):
		if not re.match("[0-9a-f]{2}([-:])[0-9a-f]{2}(\\1[0-9a-f]{2}){4}$", opts.Mac.lower()):
			print "Malformed MAC address\n"
			parser.print_help()
			exit(-1)
	
	if (opts.Change and opts.Macall):
		if not re.match("[0-9a-f]{2}([-:])[0-9a-f]{2}(\\1[0-9a-f]{2}){4}$", opts.Macall.lower()):
			print "Malformed MAC address\n"
			parser.print_help()
			exit(-1)
	
	if opts.Physicaldev:
		if devname_to_pcinum_map(opts.Physicaldev) == '':
			print "The device " + opts.Physicaldev + " doesn't exist\n"
			exit(-1)
		if not pcinum_intelsriov(devname_to_pcinum_map(opts.Physicaldev)):
			print "The device " + opts.Physicaldev + " isn't a valid SRIOV device\n"
			exit(-1)
		if not pcidev_vfcount(opts.Physicaldev):
			print "The device " + opts.Physicaldev + " doesn't have any virtual functions\n"
			exit(-1)
	
	return opts	

# Returns the pci device id from device name	
def devname_to_pcinum_map(netdev):
	runcmd = "ls -l /sys/class/net/" + netdev + "/ 2>/dev/null | grep device | awk '{ print $(NF) }' | cut -b 10-"
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	pcidevid = str.splitlines(process.communicate()[0])
	if not len(pcidevid):
		return ''
	return str.split(pcidevid[0])[0]
	
# Returns device name from pci device id
def pcinum_to_devname_map(pcidevid):
	runcmd = "ls -lR /sys/devices/pci* | grep '" + pcidevid + "' | grep net | tail -n +3 | awk -F '/' '{ print $(NF-1) }' | head -n 1"
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	netdev = str.splitlines(process.communicate()[0])
	if not len(netdev):
		return ''
	return str.split(netdev[0])[0]	

def pcinum_to_devname_map2(pcidevid):
        runcmd = "cat /tmp/pcinum_to_devname_map.cache | grep '" + pcidevid + "' | grep net | tail -n +3 | awk -F '/' '{ print $(NF-1) }' | head -n 1"
        process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
        netdev = str.splitlines(process.communicate()[0])
        if not len(netdev):
                return ''
        return str.split(netdev[0])[0]
	
# checks to see if pc device id is a compatible Intel PF
def pcinum_intelsriov(pcidevid):
	runcmd = "lspci -D | grep -E 'Intel Corporation.*(82576|82599|X540)' | grep -v 'Virtual Function' | grep " + pcidevid
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	pcidev = str.splitlines(process.communicate()[0])
	if not len(pcidev):
		return 0
	return 1

# returns a list of compatible Intel PF pcidevids
def pcinum_intelsriov_list():
	runcmd = "lspci -D | grep -E 'Intel Corporation.*(82576|82599|X540)' | grep -v 'Virtual Function' | awk '{ print $1 }'"
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	pcidev_list = str.splitlines(process.communicate()[0])
	if not len(pcidev_list):
		return ''
	return pcidev_list
	
# Returns number of vf's for a pf
def pcidev_vfcount(netdev):
	runcmd = "ls -l /sys/class/net/" + netdev + "/device/virt* 2>/dev/null"
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	netdev = str.splitlines(process.communicate()[0])
	return len(netdev)

# Returns vf pcidevid number from vf # 
def vf_to_pcidev_map(vfnum, netdev):
	runcmd = "ls -l /sys/class/net/" + netdev + "/device/virtfn" + vfnum + " | awk '{ print $11 }' | cut -b 4-"
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	pcidevid = str.splitlines(process.communicate()[0])
	if not len(pcidevid):
		exit(-1)
	return str.split(pcidevid[0])[0]

# converts pcidevid to xml
def pcidevid_to_xml(pcidevid):
	tmp = str.split(pcidevid, ':')
	domain = tmp[0]
	bus = tmp[1]
	slot = str.split(tmp[2], '.')[0]
	function = str.split(tmp[2], '.')[1]
	xmlskel = "<address domain='0x" + domain + "' bus='0x" + bus + "' slot='0x" + slot + "' function='0x" + function + "'/>"
	return xmlskel

# creates the xml file used by virsh to add/remove pci devices
def create_pci_xml_file(pcidevid):
	tmpxml = "<hostdev mode='subsystem' type='pci' managed='yes'><source>" + pcidevid_to_xml(pcidevid) + "</source></hostdev>\n"
	tmpxmlfile = open('/tmp/vf.xml', 'w')
	tmpxmlfile.write(tmpxml)
	tmpxmlfile.close()
	return 0
	
# returns the first running vm using the specified vf	
def vf_runningvm_map(pcidevid):
	runcmd = "virsh list | tail -n +3 | awk '{ print $2 }'"
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	vmlist = str.splitlines(process.communicate()[0])
	for index in range(len(vmlist)):
		if not str.split(vmlist[index]):
			return ''
		runcmd = "cat /var/run/libvirt/qemu/" + vmlist[index] + ".xml | grep \"" + pcidevid_to_xml(pcidevid) + "\""
		process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
		if str.splitlines(process.communicate()[0]):
			return vmlist[index]
	return ''

def vf_runningvm_map2(pcidevid):
        runcmd = "cat /tmp/vf_runningvm_map.cache"
        process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
        vmlist = str.splitlines(process.communicate()[0])
        for index in range(len(vmlist)):
                if not str.split(vmlist[index]):
                        return ''
                runcmd = "cat /var/run/libvirt/qemu/" + vmlist[index] + ".xml | grep \"" + pcidevid_to_xml(pcidevid) + "\""
                process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
                if str.splitlines(process.communicate()[0]):
                        return vmlist[index]
        return ''
	
# returns link status
def devname_linkstatus(netdev):
	runcmd = "ethtool " + netdev + " 2>/dev/null | grep 'Link detected' | awk '{ print $3 }'"
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	linkstatus = str.splitlines(process.communicate()[0])
	if len(linkstatus):
		return linkstatus[0]
	else:
		return ''

# returns pf driver name from devicename
def devname_driver(netdev):
	runcmd = "ethtool -i " + netdev + " | grep driver | awk '{ print $2 }'"
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	driver = str.splitlines(process.communicate()[0])
	if len(driver):
		return driver[0]
	else:
		return 'UNKNOWN'

# Returns vf info for a pf from ip utility and other functions
def vf_info_ip_utility(netdev):
	vf_List = []
	runcmd = "ip link show " + netdev + " | tail -n +3"
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	vflist = str.splitlines(process.communicate()[0])
	for index in range(len(vflist)):
		if not str.split(vflist[index]):
			return vf_List
		vf_data = VFinfo()
		tmp = str.split(vflist[index])
		vf_data.VFnumber = tmp[1]
		vf_data.VFparentdev = netdev
		vf_data.VFparentdriver = devname_driver(netdev)
		vf_data.VFpciid = vf_to_pcidev_map(vf_data.VFnumber, netdev)
		vf_data.VFflexpp = pcinum_to_devname_map2(vf_data.VFpciid)
		vf_data.VFlinkstatus = devname_linkstatus(vf_data.VFflexpp)
		vf_data.VFassignedvm = vf_runningvm_map2(vf_to_pcidev_map(vf_data.VFnumber, netdev))
		vf_data.VFmac = tmp[3].replace(',','')
		if (len(tmp) == 6):
			vf_data.VFvlan = tmp[5]
		elif (len(tmp) == 8):
			vf_data.VFtxrate = tmp[6] + tmp[7]
		elif (len(tmp) == 9):
			vf_data.VFvlan = tmp[5].replace(',','')
		elif (len(tmp) == 10):
			vf_data.VFvlan = tmp[5].replace(',','')
			vf_data.VFtxrate = tmp[8] + tmp[9]
		elif (len(tmp) == 11):
			vf_data.VFtxrate = tmp[6] + tmp[7].replace(',','')
		elif (len(tmp) == 13):
			vf_data.VFvlan = tmp[5].replace(',','')
			vf_data.VFtxrate = tmp[8] + tmp[9].replace(',','')
		
		vf_List.append(vf_data)
	return vf_List		
		
# Display vf list for a pf

def vf_display(vf_list):
	vfshowtable = prettytable.PrettyTable(["VF Number", "VF Parent Device", "VF PCI ID", "VF MAC Address", "VF VLAN", "VF TxRATE", "VF Netdevice", "VF Link Enabled", "Assigned VM"])
	if not range(len(vf_list)):
		return 0
	for index in range(len(vf_list)):
		vfshowtable.add_row([vf_list[index].VFnumber, vf_list[index].VFparentdev + '[' + vf_list[index].VFparentdriver + ']', vf_list[index].VFpciid, vf_list[index].VFmac, vf_list[index].VFvlan, vf_list[index].VFtxrate, vf_list[index].VFflexpp, vf_list[index].VFlinkstatus, vf_list[index].VFassignedvm])
	print vfshowtable
	return 0

# Dump udev rules for VFs
def vf_udev(vf_list):
	if not range(len(vf_list)):
		print "There are no VFs for this device\n" 
		return 0
	for index in range(len(vf_list)):
		if vf_list[index].VFparentdriver == 'igb': 
			print "SUBSYSTEM==\"net\", ACTION==\"add\", DRIVERS==\"igbvf\", ATTR{type}==\"1\", KERNELS==\"" + vf_list[index].VFpciid + "\", NAME=\"" + vf_list[index].VFparentdev + "_" + vf_list[index].VFnumber +"\""
		else:
			print "SUBSYSTEM==\"net\", ACTION==\"add\", DRIVERS==\"ixgbevf\", ATTR{type}==\"1\", KERNELS==\"" + vf_list[index].VFpciid + "\", NAME=\"" + vf_list[index].VFparentdev + "_" + vf_list[index].VFnumber +"\""
	return 0

# Attach a VF to a running VM
def attach_vf(vm_name, vf_num, netdev):
	create_pci_xml_file(vf_to_pcidev_map(vf_num, netdev))
	runcmd = 'virsh attach-device ' + vm_name + ' /tmp/vf.xml'
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	time.sleep(1)
	if process.poll():
		return -1
	return 0	
	
# Detach a VF from a running VM
def detach_vf(vm_name, vf_num, netdev):
	create_pci_xml_file(vf_to_pcidev_map(vf_num, netdev))
	runcmd = 'virsh detach-device ' + vm_name + ' /tmp/vf.xml'
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	time.sleep(1)
	if process.poll():
		return -1
	return 0

# Change TX rate for a VF
def vf_txrate(vf_num, netdev, tx_rate):
	if devname_driver(netdev) == 'igb' and int(tx_rate) > 1000:
			print "The igb driver doesn't support a greater tx rate than 1000Mbits, so setting rate to 1000Mbits\n"
			tx_rate = 1000
	runcmd = 'ip link set ' + netdev + ' vf ' + vf_num + ' rate ' + tx_rate
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	time.sleep(1)
	if process.poll():
		return -1
	return 0

# Change VLAN for a VF
def vf_vlan(vf_num, netdev, vlan):
	runcmd = 'ip link set ' + netdev + ' vf ' + vf_num + ' vlan ' + vlan
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	time.sleep(1)
	if process.poll():
		return -1
	return 0	

# Change MAC address for a VF
def vf_mac(vf_num, netdev, mac):
	runcmd = 'ip link set ' + netdev + ' vf ' + vf_num + ' mac ' + mac
	process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	time.sleep(1)
	if process.poll():
		return -1
	return 0	

# Change all MAC addresses for a set of VF starting from a given address
def vf_mac_all(netdev, mac):
	mac_change = mac[15:]
	mac_keep = mac[:-2]
	for index in range(pcidev_vfcount(netdev)):
		vf_mac(str(index), netdev, mac_keep + hex((int(mac_change, 16) + index + 17)%256).lstrip("0x"))
	return 0

# Displays all devices that support VFs	
def vf_display_all():	
	pci_numlist = pcinum_intelsriov_list()
	if not len(pci_numlist):
		print "No VF supported devices found\n"
		return 0
	for index in range(len(pci_numlist)):
		vf_display(vf_info_ip_utility(pcinum_to_devname_map2(pci_numlist[index])))
	return 0

# Run common commands for parsing and save them in temp
def cache_commands():
	runcmd = "ls -lR /sys/devices/pci* > /tmp/pcinum_to_devname_map.cache"
        process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	runcmd = "virsh list | tail -n +3 | awk '{ print $2 }' > /tmp/vf_runningvm_map.cache"
        process = subprocess.Popen([runcmd], shell=True, stdout=subprocess.PIPE)
	return 0

# Main Thread

if __name__=="__main__":
	options_list = parse_options()
	if options_list.Attach == True:
		attach_vf(options_list.Instance, options_list.Function, options_list.Physicaldev)
	elif options_list.Detach == True:
		detach_vf(options_list.Instance, options_list.Function, options_list.Physicaldev)
	elif options_list.Udev == True:
		vf_udev(vf_info_ip_utility(options_list.Physicaldev))
	elif options_list.Change == True:
		if options_list.Macall:
			vf_mac_all(options_list.Physicaldev, options_list.Macall)
			sys.exit(0)
		if options_list.Mac:
			vf_mac(options_list.Function, options_list.Physicaldev, options_list.Mac)
		if options_list.Rate:
			vf_txrate(options_list.Function, options_list.Physicaldev, options_list.Rate)
		if options_list.Vlan:
			vf_vlan(options_list.Function, options_list.Physicaldev, options_list.Vlan)	
	elif options_list.Physicaldev:
		cache_commands()
		vf_display(vf_info_ip_utility(options_list.Physicaldev))
	else:
		cache_commands()
		vf_display_all()
	
	sys.exit(0)
